{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "4a6ab9a2-28a2-445d-8512-a0dc8d1b54e9",
   "metadata": {},
   "source": [
    "# Power Coder\n",
    "\n",
    "1. Convert code between two programming language; supporting languages are Python, Java, JavaScript, TypeScript, C, C++, C#, Go, Rust, Kotlin, Swift, PHP, Julia\n",
    "2. Automatically add docstring/comments based on selected comment style\n",
    "3. Automatically generate unit tests based on selected unit test style\n",
    "4. Supporting models: gpt-4o, claude-3-5-sonnet-20240620, gemini-2.5-flash\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e610bf56-a46e-4aff-8de1-ab49d62b1ad3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# imports\n",
    "\n",
    "import os\n",
    "import io\n",
    "import sys\n",
    "import json\n",
    "import requests\n",
    "from dotenv import load_dotenv\n",
    "from openai import OpenAI\n",
    "import google.generativeai\n",
    "import anthropic\n",
    "from IPython.display import Markdown, display, update_display\n",
    "import gradio as gr\n",
    "import subprocess"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4f672e1c-87e9-4865-b760-370fa605e614",
   "metadata": {},
   "outputs": [],
   "source": [
    "# environment\n",
    "\n",
    "load_dotenv(override=True)\n",
    "os.environ['OPENAI_API_KEY'] = os.getenv('OPENAI_API_KEY', 'your-key-if-not-using-env')\n",
    "os.environ['ANTHROPIC_API_KEY'] = os.getenv('ANTHROPIC_API_KEY', 'your-key-if-not-using-env')\n",
    "os.environ['GOOGLE_API_KEY'] = os.getenv('GOOGLE_API_KEY', 'your-key-if-not-using-env')\n",
    "os.environ['HF_TOKEN'] = os.getenv('HF_TOKEN', 'your-key-if-not-using-env')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8aa149ed-9298-4d69-8fe2-8f5de0f667da",
   "metadata": {},
   "outputs": [],
   "source": [
    "# initialize\n",
    "\n",
    "openai = OpenAI()\n",
    "claude = anthropic.Anthropic()\n",
    "gemini_via_openai_client = OpenAI(\n",
    "    api_key=os.environ['GOOGLE_API_KEY'], \n",
    "    base_url=\"https://generativelanguage.googleapis.com/v1beta/openai/\"\n",
    ")\n",
    "OPENAI_MODEL = \"gpt-4o\"\n",
    "CLAUDE_MODEL = \"claude-3-5-sonnet-20240620\"\n",
    "GEMINI_MODEL = \"gemini-2.5-flash\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "37b204dd-f770-41d9-9b19-7e1baa5273cd",
   "metadata": {},
   "source": [
    "## 1. Convesion Part"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6896636f-923e-4a2c-9d6c-fac07828a201",
   "metadata": {},
   "outputs": [],
   "source": [
    "def convert_system_prompt_for(in_lang, out_lang):\n",
    "    convert_system_message = f\"You are an assistant that reimplements {in_lang} code in high performance {out_lang}. \"\n",
    "    convert_system_message += f\"Respond only with {out_lang} code; use comments sparingly and do not provide any explanation other than occasional comments. \"\n",
    "    convert_system_message += f\"The {out_lang} response needs to produce an identical output in the fastest possible time. Keep implementations of random number generators identical so that results match exactly.\"\n",
    "    return convert_system_message"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8e7b3546-57aa-4c29-bc5d-f211970d04eb",
   "metadata": {},
   "outputs": [],
   "source": [
    "def convert_user_prompt_for(in_lang, out_lang, input_instruct, in_code):\n",
    "    convert_user_prompt = f\"Rewrite this {in_lang} code in {out_lang} with the fastest possible implementation that produces identical output in the least time. \"\n",
    "    convert_user_prompt += f\"Respond only with {out_lang} code; do not explain your work other than a few comments. \"\n",
    "    convert_user_prompt += f\"Pay attention to number types to ensure no int overflows. Remember to include all necessary {out_lang} packages or modules, for example, iomanip for C++.\\n\\n\"\n",
    "    if input_instruct:\n",
    "        convert_user_prompt += \"Addtional instruction is: \" + input_instruct\n",
    "    convert_user_prompt += in_code\n",
    "    return convert_user_prompt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c6190659-f54c-4951-bef4-4960f8e51cc4",
   "metadata": {},
   "outputs": [],
   "source": [
    "def convert_messages_for(in_lang, out_lang, input_instruct, in_code):\n",
    "    return [\n",
    "        {\"role\": \"system\", \"content\": convert_system_prompt_for(in_lang, out_lang)},\n",
    "        {\"role\": \"user\", \"content\": convert_user_prompt_for(in_lang, out_lang, input_instruct, in_code)}\n",
    "    ]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c3b497b3-f569-420e-b92e-fb0f49957ce0",
   "metadata": {},
   "outputs": [],
   "source": [
    "python_hard = \"\"\"# Be careful to support large number sizes\n",
    "\n",
    "def lcg(seed, a=1664525, c=1013904223, m=2**32):\n",
    "    value = seed\n",
    "    while True:\n",
    "        value = (a * value + c) % m\n",
    "        yield value\n",
    "        \n",
    "def max_subarray_sum(n, seed, min_val, max_val):\n",
    "    lcg_gen = lcg(seed)\n",
    "    random_numbers = [next(lcg_gen) % (max_val - min_val + 1) + min_val for _ in range(n)]\n",
    "    max_sum = float('-inf')\n",
    "    for i in range(n):\n",
    "        current_sum = 0\n",
    "        for j in range(i, n):\n",
    "            current_sum += random_numbers[j]\n",
    "            if current_sum > max_sum:\n",
    "                max_sum = current_sum\n",
    "    return max_sum\n",
    "\n",
    "def total_max_subarray_sum(n, initial_seed, min_val, max_val):\n",
    "    total_sum = 0\n",
    "    lcg_gen = lcg(initial_seed)\n",
    "    for _ in range(20):\n",
    "        seed = next(lcg_gen)\n",
    "        total_sum += max_subarray_sum(n, seed, min_val, max_val)\n",
    "    return total_sum\n",
    "\n",
    "# Parameters\n",
    "n = 10000         # Number of random numbers\n",
    "initial_seed = 42 # Initial seed for the LCG\n",
    "min_val = -10     # Minimum value of random numbers\n",
    "max_val = 10      # Maximum value of random numbers\n",
    "\n",
    "# Timing the function\n",
    "import time\n",
    "start_time = time.time()\n",
    "result = total_max_subarray_sum(n, initial_seed, min_val, max_val)\n",
    "end_time = time.time()\n",
    "\n",
    "print(\"Total Maximum Subarray Sum (20 runs):\", result)\n",
    "print(\"Execution Time: {:.6f} seconds\".format(end_time - start_time))\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0be9f47d-5213-4700-b0e2-d444c7c738c0",
   "metadata": {},
   "outputs": [],
   "source": [
    "def convert_stream_gpt(in_lang, out_lang, input_instruct, in_code):    \n",
    "    stream = openai.chat.completions.create(model=OPENAI_MODEL, messages=convert_messages_for(in_lang, out_lang, input_instruct, in_code), temperature=0.0, stream=True)\n",
    "    reply = \"\"\n",
    "    for chunk in stream:\n",
    "        fragment = chunk.choices[0].delta.content or \"\"\n",
    "        reply += fragment\n",
    "        yield reply"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8669f56b-8314-4582-a167-78842caea131",
   "metadata": {},
   "outputs": [],
   "source": [
    "def convert_stream_claude(in_lang, out_lang, input_instruct, in_code):\n",
    "    result = claude.messages.stream(\n",
    "        model=CLAUDE_MODEL,\n",
    "        max_tokens=2000,\n",
    "        temperature=0.0,\n",
    "        system=convert_system_prompt_for(in_lang, out_lang),\n",
    "        messages=[{\"role\": \"user\", \"content\": convert_user_prompt_for(in_lang, out_lang, input_instruct, in_code)}],\n",
    "    )\n",
    "    reply = \"\"\n",
    "    with result as stream:\n",
    "        for text in stream.text_stream:\n",
    "            reply += text\n",
    "            yield reply"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "01d3cd4f-c100-4e25-8670-0663513f6136",
   "metadata": {},
   "outputs": [],
   "source": [
    "def convert_stream_gemini(in_lang, out_lang, input_instruct, in_code):    \n",
    "    stream = gemini_via_openai_client.chat.completions.create(model=GEMINI_MODEL, messages=convert_messages_for(in_lang, out_lang, input_instruct, in_code), temperature=0.0, stream=True)\n",
    "    reply = \"\"\n",
    "    for chunk in stream:\n",
    "        fragment = chunk.choices[0].delta.content or \"\"\n",
    "        reply += fragment\n",
    "        yield reply"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2f1ae8f5-16c8-40a0-aa18-63b617df078d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def optimize(in_lang, out_lang, in_code, input_instruct, convert_model):\n",
    "    if \"gpt\" in convert_model.lower():\n",
    "        result = convert_stream_gpt(in_lang, out_lang, input_instruct, in_code)\n",
    "    elif \"claude\" in convert_model.lower():\n",
    "        result = convert_stream_claude(in_lang, out_lang, input_instruct, in_code)\n",
    "    elif \"gemini\" in convert_model.lower():\n",
    "        result = convert_stream_gemini(in_lang, out_lang, input_instruct, in_code)\n",
    "    else:\n",
    "        raise ValueError(\"Unknown convert model\")\n",
    "    for stream_so_far in result:\n",
    "        yield stream_so_far        "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "07383878-f887-464f-8bc7-527c669d3edd",
   "metadata": {},
   "source": [
    "## 2. Comment part"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d254038c-fdd6-4ef8-8b7a-a074f1e7405d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def comment_system_prompt_for(lang, comment_style):\n",
    "    comment_system_message = f\"You are an assistant that generate necessary, concise and clear comment/docstring for the {lang} code by applying {comment_style} comment style. \"\n",
    "    comment_system_message += f\"Respond only with added comments, and do not provide any redundant explanation. \"\n",
    "    return comment_system_message"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e95cee4f-f229-4c9f-8e67-8a68cc9534c3",
   "metadata": {},
   "outputs": [],
   "source": [
    "def comment_user_prompt_for(lang, code, comment_style):\n",
    "    comment_user_prompt = f\"Add the comments/docstring on the given code for the {lang} programming language in {comment_style} comment style. \"\n",
    "    comment_user_prompt += f\"Respond only with added comments, and do not provide any redundant explanation.\\n\\n\"\n",
    "    comment_user_prompt += f\"The given code is as follows: \"\n",
    "    comment_user_prompt += code\n",
    "    return comment_user_prompt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "507426c2-cf5a-4041-b904-b18a5afe83b6",
   "metadata": {},
   "outputs": [],
   "source": [
    "def comment_messages_for(lang, code, comment_style):\n",
    "    return [\n",
    "        {\"role\": \"system\", \"content\": comment_system_prompt_for(lang, comment_style)},\n",
    "        {\"role\": \"user\", \"content\": comment_user_prompt_for(lang, code, comment_style)}\n",
    "    ]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7e1c8cf6-7a15-4e79-82f6-6bb2a0b85773",
   "metadata": {},
   "outputs": [],
   "source": [
    "def comment_stream_gpt(lang, code, comment_style):    \n",
    "    stream = openai.chat.completions.create(model=OPENAI_MODEL, messages=comment_messages_for(lang, code, comment_style), temperature=0.0, stream=True)\n",
    "    reply = \"\"\n",
    "    for chunk in stream:\n",
    "        fragment = chunk.choices[0].delta.content or \"\"\n",
    "        reply += fragment\n",
    "        yield reply"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "26f27781-4a3e-4e5f-a8ab-9a25944a9879",
   "metadata": {},
   "outputs": [],
   "source": [
    "def comment_stream_claude(lang, code, comment_style):\n",
    "    result = claude.messages.stream(\n",
    "        model=CLAUDE_MODEL,\n",
    "        max_tokens=2000,\n",
    "        temperature=0.0,\n",
    "        system=comment_system_prompt_for(lang, comment_style),\n",
    "        messages=[{\"role\": \"user\", \"content\": comment_user_prompt_for(lang, code, comment_style)}],\n",
    "    )\n",
    "    reply = \"\"\n",
    "    with result as stream:\n",
    "        for text in stream.text_stream:\n",
    "            reply += text\n",
    "            yield reply"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8e6719e7-f2f3-40ea-8fed-01d84a641306",
   "metadata": {},
   "outputs": [],
   "source": [
    "def comment_stream_gemini(lang, code, comment_style):    \n",
    "    stream = gemini_via_openai_client.chat.completions.create(model=GEMINI_MODEL, messages=comment_messages_for(lang, code, comment_style), temperature=0.0, stream=True)\n",
    "    reply = \"\"\n",
    "    for chunk in stream:\n",
    "        fragment = chunk.choices[0].delta.content or \"\"\n",
    "        reply += fragment\n",
    "        yield reply"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2b98acc4-23d8-4671-8f19-92d72631b55d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_comments_via_model(lang, code, comment_style, comment_model):\n",
    "    if \"gpt\" in comment_model.lower():\n",
    "        result = comment_stream_gpt(lang, code, comment_style)\n",
    "    elif \"claude\" in comment_model.lower():\n",
    "        result = comment_stream_claude(lang, code, comment_style)\n",
    "    elif \"gemini\" in comment_model.lower():\n",
    "        result = comment_stream_gemini(lang, code, comment_style)\n",
    "    else:\n",
    "        raise ValueError(\"Unknown comment model\")\n",
    "    for stream_so_far in result:\n",
    "        yield stream_so_far        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "282c75ae-d8c3-4866-a024-f7ecf87b3cde",
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_comments_fn(comment_option, in_lang, out_lang, in_code, out_code, in_comment_style, out_comment_style, comment_model):\n",
    "    if 'input' in comment_option:\n",
    "        in_gen = generate_comments_via_model(in_lang, in_code, in_comment_style, comment_model)\n",
    "        for in_output in in_gen:\n",
    "            yield in_output, \"\"\n",
    "    elif 'output' in comment_option:\n",
    "        out_gen = generate_comments_via_model(out_lang, out_code, out_comment_style, comment_model)\n",
    "        for out_output in out_gen:\n",
    "            yield \"\", out_output\n",
    "    elif 'both' in comment_option:\n",
    "        in_gen = generate_comments_via_model(in_lang, in_code, in_comment_style, comment_model)\n",
    "        out_gen = generate_comments_via_model(out_lang, out_code, out_comment_style, comment_model)\n",
    "        for in_output, out_output in zip(in_gen, out_gen):\n",
    "            yield in_output, out_output"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ce2c178c-d03c-49c0-b0e9-c57c699bca08",
   "metadata": {},
   "source": [
    "## 3. Unit test part"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e5a4743e-e1a8-42c7-8f1f-a73d49c0895d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def unit_test_system_prompt_for(lang, unit_test_style):\n",
    "    unit_test_system_message = f\"You are an assistant that generate necessary, concise, clear and executable unit tests for the {lang} code by applying {unit_test_style} unit test style. \"\n",
    "    unit_test_system_message += f\"Respond only with generated unit tests; use comments sparingly and do not provide any explanation other than occasional comments. \"\n",
    "    return unit_test_system_message"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "334d5e40-71ff-4d24-8cef-b6c81c188e4d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def unit_test_user_prompt_for(lang, code, unit_test_style):\n",
    "    unit_test_user_prompt = f\"Add the unit tests on the given code for the {lang} programming language in {unit_test_style} unit test style. \"\n",
    "    unit_test_user_prompt += f\"Respond only with generated unit tests; use comments sparingly and do not provide any explanation other than occasional comments.\\n\\n\"\n",
    "    unit_test_user_prompt += f\"The given code is as follows: \"\n",
    "    unit_test_user_prompt += code\n",
    "    return unit_test_user_prompt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8a8e061f-3993-4746-9425-d938d2537f65",
   "metadata": {},
   "outputs": [],
   "source": [
    "def unit_test_messages_for(lang, code, unit_test_style):\n",
    "    return [\n",
    "        {\"role\": \"system\", \"content\": unit_test_system_prompt_for(lang, unit_test_style)},\n",
    "        {\"role\": \"user\", \"content\": unit_test_user_prompt_for(lang, code, unit_test_style)}\n",
    "    ]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "71c1613b-7a16-4443-acec-d0a2d9bed192",
   "metadata": {},
   "outputs": [],
   "source": [
    "def unit_test_stream_gpt(lang, code, unit_test_style):    \n",
    "    stream = openai.chat.completions.create(model=OPENAI_MODEL, messages=unit_test_messages_for(lang, code, unit_test_style), stream=True)\n",
    "    reply = \"\"\n",
    "    for chunk in stream:\n",
    "        fragment = chunk.choices[0].delta.content or \"\"\n",
    "        reply += fragment\n",
    "        yield reply"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8a6e3502-f7ff-42b8-8fc5-2697b2d1f36e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def unit_test_stream_claude(lang, code, unit_test_style):\n",
    "    result = claude.messages.stream(\n",
    "        model=CLAUDE_MODEL,\n",
    "        max_tokens=2000,\n",
    "        system=unit_test_system_prompt_for(lang, unit_test_style),\n",
    "        messages=[{\"role\": \"user\", \"content\": unit_test_user_prompt_for(lang, code, unit_test_style)}],\n",
    "    )\n",
    "    reply = \"\"\n",
    "    with result as stream:\n",
    "        for text in stream.text_stream:\n",
    "            reply += text\n",
    "            yield reply"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8d7f694f-a276-4bdc-9cfb-755483fd4380",
   "metadata": {},
   "outputs": [],
   "source": [
    "def unit_test_stream_gemini(lang, code, unit_test_style):    \n",
    "    stream = gemini_via_openai_client.chat.completions.create(model=GEMINI_MODEL, messages=unit_test_messages_for(lang, code, unit_test_style), stream=True)\n",
    "    reply = \"\"\n",
    "    for chunk in stream:\n",
    "        fragment = chunk.choices[0].delta.content or \"\"\n",
    "        reply += fragment\n",
    "        yield reply"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c824429a-b18a-4320-8258-0141037a6531",
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_unit_test_via_model(lang, code, unit_test_style, unit_test_model):\n",
    "    if \"gpt\" in unit_test_model.lower():\n",
    "        result = unit_test_stream_gpt(lang, code, unit_test_style)\n",
    "    elif \"claude\" in unit_test_model.lower():\n",
    "        result = unit_test_stream_claude(lang, code, unit_test_style)\n",
    "    elif \"gemini\" in unit_test_model.lower():\n",
    "        result = unit_test_stream_gemini(lang, code, unit_test_style)\n",
    "    else:\n",
    "        raise ValueError(\"Unknown unit test model\")\n",
    "    for stream_so_far in result:\n",
    "        yield stream_so_far        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c3e59e26-37c0-4429-b69c-deb581423dd0",
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_unit_test_fn(unit_test_option, in_lang, out_lang, in_code, out_code, in_unit_test_style, out_unit_test_style, unit_test_model):\n",
    "    if 'input' in unit_test_option:\n",
    "        in_gen = generate_unit_test_via_model(in_lang, in_code, in_unit_test_style, unit_test_model)\n",
    "        for in_output in in_gen:\n",
    "            yield in_output, \"\"\n",
    "    elif 'output' in unit_test_option:\n",
    "        out_gen = generate_unit_test_via_model(out_lang, out_code, out_unit_test_style, unit_test_model)\n",
    "        for out_output in out_gen:\n",
    "            yield \"\", out_output\n",
    "    elif 'both' in unit_test_option:\n",
    "        in_gen = generate_unit_test_via_model(in_lang, in_code, in_unit_test_style, unit_test_model)\n",
    "        out_gen = generate_unit_test_via_model(out_lang, out_code, out_unit_test_style, unit_test_model)\n",
    "        for in_output, out_output in zip(in_gen, out_gen):\n",
    "            yield in_output, out_output"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2a1f4d0c-f417-4de4-be9f-441cbe5a6db3",
   "metadata": {},
   "source": [
    "## 4. Gradio UI part"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9a2274f1-d03b-42c0-8dcc-4ce159b18442",
   "metadata": {},
   "outputs": [],
   "source": [
    "LANGUAGE_INFO = {\n",
    "    \"Python\": {\n",
    "        \"doc_style\": [\"Google-style\", \"NumPy-style\", \"reST\", \"Doxygen\"],\n",
    "        \"unit_test_style\": [\"unittest\", \"pytest\", \"doctest\"]\n",
    "    },\n",
    "    \"Java\": {\n",
    "        \"doc_style\": [\"Javadoc\"],\n",
    "        \"unit_test_style\": [\"JUnit4\", \"JUnit5\", \"TestNG\"]\n",
    "    },\n",
    "    \"JavaScript\": {\n",
    "        \"doc_style\": [\"JSDoc\"],\n",
    "        \"unit_test_style\": [\"Jest\", \"Mocha + Chai\", \"Jasmine\"]\n",
    "    },\n",
    "    \"TypeScript\": {\n",
    "        \"doc_style\": [\"JSDoc\", \"TSDoc\"],\n",
    "        \"unit_test_style\": [\"Jest\", \"Mocha + Chai\", \"Vitest\"]\n",
    "    },\n",
    "    \"C\": {\n",
    "        \"doc_style\": [\"Doxygen\"],\n",
    "        \"unit_test_style\": [\"Google Test (gtest)\", \"CppUnit\", \"Catch2\"]\n",
    "    },\n",
    "    \"C++\": {\n",
    "        \"doc_style\": [\"Doxygen\"],\n",
    "        \"unit_test_style\": [\"Google Test (gtest)\", \"CppUnit\", \"Catch2\"]\n",
    "    },\n",
    "    \"C#\": {\n",
    "        \"doc_style\": [\"XML comments\"],\n",
    "        \"unit_test_style\": [\"xUnit\", \"NUnit\", \"MSTest\"]\n",
    "    },\n",
    "    \"Go\": {\n",
    "        \"doc_style\": [\"Godoc\"],\n",
    "        \"unit_test_style\": [\"Built-in testing package\"]\n",
    "    },\n",
    "    \"Rust\": {\n",
    "        \"doc_style\": [\"Rustdoc\", \"Markdown\"],\n",
    "        \"unit_test_style\": [\"Built-in #[test] annotation\"]\n",
    "    },\n",
    "    \"Kotlin\": {\n",
    "        \"doc_style\": [\"KDoc\"],\n",
    "        \"unit_test_style\": [\"JUnit\", \"Kotest\", \"Spek\"]\n",
    "    },\n",
    "    \"Swift\": {\n",
    "        \"doc_style\": [\"Mark-style comments\"],\n",
    "        \"unit_test_style\": [\"XCTest\"]\n",
    "    },\n",
    "    \"PHP\": {\n",
    "        \"doc_style\": [\"PHPDoc\"],\n",
    "        \"unit_test_style\": [\"PHPUnit\"]\n",
    "    },\n",
    "    \"Julia\": {\n",
    "        \"doc_style\": [\"Markdown\"],\n",
    "        \"unit_test_style\": [\"Built-in Test standard library\"]\n",
    "    }\n",
    "}\n",
    "LANGUAGES = list(LANGUAGE_INFO.keys())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b50e7833-8f6f-407e-8174-37af9cec2030",
   "metadata": {},
   "outputs": [],
   "source": [
    "with gr.Blocks(title=\"Power Coder\", theme=gr.themes.Citrus(), css=\"\"\"\n",
    ".selected {\n",
    "    background-color: orange !important;\n",
    "    box-shadow: 0 4px 12px rgba(255, 140, 0, 0.5) !important;\n",
    "    color: black;\n",
    "}\n",
    ".unselected {\n",
    "    background-color: gray !important;\n",
    "    box-shadow: 0 4px 12px rgba(128, 128, 128, 0.4);\n",
    "    color: white;\n",
    "}\n",
    "\"\"\") as ui:\n",
    "    current_selected = gr.State(\"\")\n",
    "    initial_in_lang = \"Python\"\n",
    "    initial_out_lang = \"Java\"\n",
    "    in_comment_style_choices = [\"Standard\"] + LANGUAGE_INFO[initial_in_lang][\"doc_style\"]\n",
    "    out_comment_style_choices = [\"Standard\"] + LANGUAGE_INFO[initial_out_lang][\"doc_style\"]\n",
    "    in_unit_test_style_choices = [\"Standard\"] + LANGUAGE_INFO[initial_in_lang][\"unit_test_style\"]\n",
    "    out_unit_test_style_choices = [\"Standard\"] + LANGUAGE_INFO[initial_out_lang][\"unit_test_style\"]\n",
    "    in_code_file_name = gr.State(\"in_code.txt\")\n",
    "    out_code_file_name = gr.State(\"out_code.txt\")\n",
    "    in_comments_file_name = gr.State(\"in_comments.txt\")\n",
    "    out_comments_file_name = gr.State(\"out_comments.txt\")\n",
    "    in_unit_test_file_name = gr.State(\"in_unit_tests.txt\")\n",
    "    out_unit_test_file_name = gr.State(\"out_unit_tests.txt\")\n",
    "    \n",
    "    \n",
    "    gr.Markdown(\"## Code Helper\")\n",
    "\n",
    "    def load_file_content(file):\n",
    "        if file is None:\n",
    "            return \"\"\n",
    "        with open(file.name, \"r\", encoding=\"utf-8\") as f:\n",
    "            return f.read()\n",
    "\n",
    "    def change_lang(lang):\n",
    "        comment_style_choices = [\"Standard\"] + LANGUAGE_INFO[lang][\"doc_style\"]\n",
    "        unit_test_style_choices = [\"Standard\"] + LANGUAGE_INFO[lang][\"unit_test_style\"]\n",
    "        return (\n",
    "            gr.update(choices=comment_style_choices, value=str(comment_style_choices[0])), \n",
    "            gr.update(choices=unit_test_style_choices, value=str(unit_test_style_choices[0]))\n",
    "        )\n",
    "\n",
    "    def download_fn(in_text, out_text, in_file_name, out_file_name):\n",
    "        if in_text:\n",
    "            with open(in_file_name, \"w\") as f:\n",
    "                f.write(in_text)\n",
    "        if out_text:\n",
    "            with open(out_file_name, \"w\") as f:\n",
    "                f.write(out_text)\n",
    "        \n",
    "    # Conversion part\n",
    "    with gr.Row():\n",
    "        in_lang = gr.Dropdown(choices=LANGUAGES, label=\"Select input language\", value=initial_in_lang, interactive=True)\n",
    "        out_lang = gr.Dropdown(choices=LANGUAGES, label=\"Select output language\", value=initial_out_lang, interactive=True)\n",
    "    with gr.Row():\n",
    "        input_file = gr.File(label=\"Upload a source code file or input below\")\n",
    "        input_instruct = gr.Textbox(\n",
    "            label=\"Additional instruction(optional)\",\n",
    "            placeholder=\"Enter the instruction you want the ouput code to follow...\\n\\nFor example: Define the variable using snake_case style.\",\n",
    "            lines=8\n",
    "        )\n",
    "    with gr.Row():\n",
    "        in_code = gr.Textbox(label=\"Input Code:\", value=python_hard, lines=10)\n",
    "        out_code = gr.Textbox(label=\"Output Code:\", lines=10)\n",
    "    with gr.Row():\n",
    "        convert_model = gr.Dropdown([\"Claude\", \"GPT\", \"Gemini\"], label=\"Select model\", value=\"Claude\")\n",
    "    with gr.Row():\n",
    "        convert = gr.Button(\"Convert code\")\n",
    "        download_code = gr.Button(\"Download code\")\n",
    "\n",
    "    gr.HTML(\"<hr style='border: none; height: 1px; background-color: #333;'>\")\n",
    "\n",
    "    def show_comment(current_selected):\n",
    "        if current_selected == \"comment\":\n",
    "            return (\n",
    "                gr.update(visible=False),\n",
    "                gr.update(visible=False),\n",
    "                gr.update(elem_classes=[\"unselected\"]),\n",
    "                gr.update(elem_classes=[\"unselected\"]),\n",
    "                \"\"\n",
    "            )\n",
    "        else:\n",
    "            return (\n",
    "                gr.update(visible=True),\n",
    "                gr.update(visible=False),\n",
    "                gr.update(elem_classes=[\"selected\"]),\n",
    "                gr.update(elem_classes=[\"unselected\"]),\n",
    "                \"comment\"\n",
    "            )\n",
    "\n",
    "    def show_unit_test(current_selected):\n",
    "        if current_selected == \"unit_test\":\n",
    "            return (\n",
    "                gr.update(visible=False),\n",
    "                gr.update(visible=False),\n",
    "                gr.update(elem_classes=[\"unselected\"]),\n",
    "                gr.update(elem_classes=[\"unselected\"]),\n",
    "                \"\"\n",
    "            )\n",
    "        else:\n",
    "            return (\n",
    "                gr.update(visible=False),\n",
    "                gr.update(visible=True),\n",
    "                gr.update(elem_classes=[\"unselected\"]),\n",
    "                gr.update(elem_classes=[\"selected\"]),\n",
    "                \"unit_test\"\n",
    "            )\n",
    "    \n",
    "    with gr.Blocks() as demo:\n",
    "        with gr.Row():\n",
    "            comment_show_up = gr.Button(\"Comment\", elem_id=\"comment-btn\", elem_classes=[\"unselected\"])\n",
    "            unit_test_show_up = gr.Button(\"Unit Test\",  elem_id=\"unit-test-btn\", elem_classes=[\"unselected\"])\n",
    "        \n",
    "        comment_section = gr.Column(visible=False)\n",
    "        unit_test_section = gr.Column(visible=False)\n",
    "        \n",
    "    with comment_section:\n",
    "        # Comment section\n",
    "        with gr.Row():\n",
    "            comment_option = gr.Radio(\n",
    "                choices=[\n",
    "                    \"Comment input code\",\n",
    "                    \"Comment output code\",\n",
    "                    \"Comment both\"\n",
    "                ],\n",
    "                label=\"Commenting Options\",\n",
    "                value=\"Comment input code\",\n",
    "                interactive=True\n",
    "            )\n",
    "        with gr.Row():\n",
    "            in_comment_style = gr.Dropdown(choices=in_comment_style_choices, label=\"Select comment style for input code\", value=in_comment_style_choices[0], interactive=True)\n",
    "            out_comment_style = gr.Dropdown(choices=out_comment_style_choices, label=\"Select comment style for oupt code\", value=out_comment_style_choices[0], interactive=True)\n",
    "        with gr.Row():\n",
    "            comment_model = gr.Dropdown([\"Claude\", \"GPT\", \"Gemini\"], label=\"Select model\", value=\"Claude\")\n",
    "        with gr.Row():\n",
    "            generate_comments = gr.Button(\"Generate comments\")\n",
    "            download_comments = gr.Button(\"Download comments\")\n",
    "        with gr.Row():\n",
    "            in_comments = gr.Textbox(label=\"Comments for Input Code:\", lines=10)\n",
    "            out_comments = gr.Textbox(label=\"Comments for Output Code:\", lines=10)\n",
    "    \n",
    "    with unit_test_section:\n",
    "        # Unit test part\n",
    "        with gr.Row():\n",
    "            unit_test_option = gr.Radio(\n",
    "                choices=[\n",
    "                    \"Add unit test for input code\",\n",
    "                    \"Add unit test for output code\",\n",
    "                    \"Add unit test for both\"\n",
    "                ],\n",
    "                label=\"Unit Test Options\",\n",
    "                value=\"Add unit test for input code\",\n",
    "                interactive=True\n",
    "            )\n",
    "        with gr.Row():\n",
    "            in_unit_test_style = gr.Dropdown(choices=in_unit_test_style_choices, label=\"Select unit test style for input code\", value=in_unit_test_style_choices[0], interactive=True)\n",
    "            out_unit_test_style = gr.Dropdown(choices=out_unit_test_style_choices, label=\"Select unit test style for oupt code\", value=out_unit_test_style_choices[0], interactive=True)\n",
    "        with gr.Row():\n",
    "            unit_test_model = gr.Dropdown([\"Claude\", \"GPT\", \"Gemini\"], label=\"Select model\", value=\"Claude\")\n",
    "        with gr.Row():\n",
    "            generate_unit_test = gr.Button(\"Generate unit test\")\n",
    "            download_unit_test = gr.Button(\"Download unit text\")\n",
    "        with gr.Row():\n",
    "            in_unit_test = gr.Textbox(label=\"Unit Test for Input Code:\", lines=10)\n",
    "            out_unit_test = gr.Textbox(label=\"Unit Test for Output Code:\", lines=10)\n",
    "\n",
    "    in_lang.change(fn=change_lang, inputs=in_lang, outputs=[in_comment_style, in_unit_test_style])\n",
    "    out_lang.change(fn=change_lang, inputs=out_lang, outputs=[out_comment_style, out_unit_test_style])\n",
    "    input_file.change(fn=load_file_content, inputs=input_file, outputs=in_code)\n",
    "    \n",
    "    convert.click(optimize, inputs=[in_lang, out_lang, in_code, input_instruct, convert_model], outputs=[out_code])\n",
    "    download_code.click(download_fn, inputs=[in_code, out_code, in_code_file_name, out_code_file_name])\n",
    "    \n",
    "    comment_show_up.click(fn=show_comment, inputs=current_selected, outputs=[comment_section, unit_test_section, comment_show_up, unit_test_show_up, current_selected])\n",
    "    unit_test_show_up.click(fn=show_unit_test, inputs=current_selected, outputs=[comment_section, unit_test_section, comment_show_up, unit_test_show_up, current_selected])\n",
    "\n",
    "    generate_comments.click(generate_comments_fn, inputs=[comment_option, in_lang, out_lang, in_code, out_code, in_comment_style, out_comment_style, comment_model], outputs=[in_comments, out_comments])\n",
    "    download_comments.click(download_fn, inputs=[in_comments, out_comments, in_comments_file_name, out_comments_file_name])\n",
    "    generate_unit_test.click(generate_unit_test_fn, inputs=[unit_test_option, in_lang, out_lang, in_code, out_code, in_unit_test_style, out_unit_test_style, unit_test_model], outputs=[in_unit_test, out_unit_test])\n",
    "    download_unit_test.click(download_fn, inputs=[in_unit_test, out_unit_test, in_unit_test_file_name, out_unit_test_file_name])\n",
    "    \n",
    "ui.launch()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0266734c-0bee-46c0-9b17-9fd2ae86cc3a",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
